"""
@Author: Bright
@File: business_tour_assistant.py
@Time: 2025/10/10
@Desc: Langgraph 多工具的使用
（用户输入）-> get_weather -> get_flights->llm_decide -> (输出建议)
"""
from typing import TypedDict

from langgraph.constants import END, START
from langgraph.graph import StateGraph
from IPython.display import Image, display


# 定义sate的具体的数据结构，如果不定义，就不会传递。
class State(TypedDict):
    flights: list
    weather: str
    city: str
    suggestion: str


graph = StateGraph(State)


# 天气工具
def get_weather(state):
    city = state["city"]
    return {"weather": f"{city} 今天晴28度", **state}


# 航班工具
def get_flights(state):
    return {"flights": ["MU5123", "CA8899"], **state}


## LLM
def llm_decide(state):
    weather = state["weather"]
    flights = state["flights"]
    return {"suggestion": f"{weather},有航班{flights}，建议穿薄外套", **state}


graph.add_node("get_weather", get_weather)
graph.add_node("get_flights", get_flights)
graph.add_node("llm_decide", llm_decide)

graph.add_edge(START, "get_weather")
graph.add_edge("get_weather", "get_flights")
graph.add_edge("get_flights", "llm_decide")
graph.add_edge("llm_decide", END)
chain = graph.compile()

# 显示流程图，mermaid
display(Image(chain.get_graph().draw_mermaid_png()))
state = chain.invoke({"city": "上海"})
print(state)
